import face_model
import argparse
import cv2
import sys
import numpy as np
from scipy import spatial
from pathlib import Path
from tqdm import tqdm

extensions=[".jpg", ".jpeg", ".png"]

parser = argparse.ArgumentParser(description='face model test')
# general
parser.add_argument('--data-path', default='', help='face image path, one image should only contain one image.')
parser.add_argument('--image-size', default='112,112', help='')
parser.add_argument('--model', default='/Users/sshuair/AI/face-recognition/insightface/models/model-r34-amf/model,0', help='path to load model.')
parser.add_argument('--ga-model', default='/Users/sshuair/AI/face-recognition/insightface/models/gamodel-r50/model,0', help='path to load model.')
parser.add_argument('--gpu', default=-1, type=int, help='gpu id')
parser.add_argument('--det', default=0, type=int, help='mtcnn option, 1 means using R+O, 0 means detect from begining')
parser.add_argument('--flip', default=0, type=int, help='whether do lr flip aug')
parser.add_argument('--threshold', default=1.24, type=float, help='ver dist threshold')
args = parser.parse_args()

model = face_model.FaceModel(args)

f = open('face_feat_db.txt', 'w')
# 每个人搜集10张照片，然后用最近邻搜索，需要包括不同年龄、环境、光照、遮罩下的人脸。
files = [x for x in Path(args.data_path).glob('**/*') if x.suffix.lower() in extensions and '._' not in str(x)]
face_feat_db = []
for item in tqdm(files):
    print(item)
    name = item.parent.parts[-1]
    img = cv2.imread(str(item))
    img = model.get_input(img)
    feat = model.get_feature(img)
    f.write(name+','+','.join([str(x) for x in feat])+'\n')
f.close()

# img = cv2.imread('Tom_Hanks_54745.png')
# print(img.shape)
# # img = cv2.imread('ws-0.jpg')
# img = model.get_input(img)
# print(img.shape)
# f1 = model.get_feature(img)
# # print(f1)
# # gender, age = model.get_ga(img)
# # print(gender)
# # print(age)
# # sys.exit(0)
# # img = cv2.imread('/raid5data/dplearn/megaface/facescrubr/112x112/Tom_Hanks/Tom_Hanks_54733.png')
# # img2 = cv2.imread('tom_hanks.jpg')
# img2 = cv2.imread('ws-2.jpg')
# img2 = model.get_input(img2)
# f2 = model.get_feature(img2)
# dist = np.sum(np.square(f1-f2))
# print('distance: ', dist)
# sim = np.dot(f1, f2.T)
# print(sim)
# #diff = np.subtract(source_feature, target_feature)
# #dist = np.sum(np.square(diff),1)

# # cos dist
# cos_dist = 1 - spatial.distance.cosine(f1, f2)
# print(cos_dist)